/**
 * Copyright 2012 Brigham Young University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package edu.byu.nlp.cluster.mom;

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;

import org.apache.commons.math3.linear.ArrayRealVector;
import org.apache.commons.math3.linear.RealMatrix;
import org.apache.commons.math3.linear.RealVector;
import org.apache.commons.math3.random.RandomGenerator;

import com.google.common.base.Preconditions;

import edu.byu.nlp.cluster.Dataset;
import edu.byu.nlp.cluster.Datasets;
import edu.byu.nlp.cluster.em.ParameterInitializer;
import edu.byu.nlp.data.SparseFeatureVector;
import edu.byu.nlp.math.RealVectors;
import edu.byu.nlp.stats.DirichletDistribution;
import edu.byu.nlp.util.Collections3;
import edu.byu.nlp.util.DoubleArrays;
import edu.byu.nlp.util.Enumeration;
import edu.byu.nlp.util.Iterables2;

/**
 * @author rah67
 *
 */
public class MoMParameterInitializer implements ParameterInitializer<MoMParameters> {

	public static interface POfYInitializer {
		double[] initialize(Dataset data, int numClusters);
	}

	public static interface POfXGivenYInitializer {
		double[][] initialize(Dataset data, int numClusters);
	}
	
	private final POfYInitializer pOfYInitializer;
	private final POfXGivenYInitializer pOfXGivenYInitializer;
	
	public MoMParameterInitializer(POfYInitializer pOfYInitializer,
			POfXGivenYInitializer pOfXGivenYInitializer) {
		this.pOfYInitializer = pOfYInitializer;
		this.pOfXGivenYInitializer = pOfXGivenYInitializer;
	}

	/** {@inheritDoc} */
	@Override
	public MoMParameters initialize(Dataset data, int numClusters) {
		double[] pOfY = pOfYInitializer.initialize(data, numClusters);
		double[][] pOfXGivenY = pOfXGivenYInitializer.initialize(data, numClusters);
		return MoMParameters.fromProbabilities(pOfY, pOfXGivenY, false, false);
	}
	
	public static POfYInitializer uniformPOfY() {
		return new POfYInitializer() {

			@Override
			public double[] initialize(Dataset data, int numClusters) {
				double[] pOfY = new double[numClusters];
				Arrays.fill(pOfY, 1.0 / numClusters);
				return pOfY;
			}
			
		};
	}
	
	public static POfYInitializer samplePOfYFromUniformDirichlet(final RandomGenerator rnd) {
		return new POfYInitializer() {

			@Override
			public double[] initialize(Dataset data, int numClusters) {
				return DirichletDistribution.sample(RealVectors.constant(1.0, numClusters), rnd);
			}
			
		};
	}
	
	public static POfYInitializer sampleFromPosteriorPOfY(final double prior, final RandomGenerator rnd) {
		return new POfYInitializer() {
			
			@Override
			public double[] initialize(Dataset data, int numClusters) {
				double[] posterior = MixtureOfMultinomialsUtil.posteriorAlpha(data, prior, numClusters);
				return DirichletDistribution.sample(posterior, rnd);
			}
		};
	}
	
	public static POfXGivenYInitializer samplePOfXGivenYFromUniformDirichlet(final RandomGenerator rnd) {
		return new POfXGivenYInitializer() {
			
			@Override
			public double[][] initialize(Dataset data, int numClusters) {
				return DirichletDistribution.sample(RealVectors.constant(1.0, data.getNumFeatures()), numClusters, rnd);
			}
		};
	}

	public static POfXGivenYInitializer samplePOfXGivenYFromDirichlet(final double alpha, final RandomGenerator rnd) {
		return new POfXGivenYInitializer() {
			
			@Override
			public double[][] initialize(Dataset data, int numClusters) {
				return DirichletDistribution.sample(RealVectors.constant(alpha, data.getNumFeatures()), numClusters,
						rnd);
			}
		};
	}

	public static POfXGivenYInitializer samplePOfXGivenYFromDirichlet(final RealVector alpha,
			final RandomGenerator rnd) {
		return new POfXGivenYInitializer() {
			
			@Override
			public double[][] initialize(Dataset data, int numClusters) {
				Preconditions.checkArgument(data.getNumFeatures() == alpha.getDimension());
				return DirichletDistribution.sample(alpha, numClusters, rnd);
			}
		};
	}

	public static POfXGivenYInitializer samplePOfXGivenYFromDirichlet(final RealMatrix alpha,
			final RandomGenerator rnd) {
		return new POfXGivenYInitializer() {
			
			@Override
			public double[][] initialize(Dataset data, int numClusters) {
				Preconditions.checkArgument(alpha.getColumnDimension() == data.getNumFeatures(),
						"This initializer can only be used on datasets of the same dimensions as the prior");
				Preconditions.checkArgument(alpha.getRowDimension() == numClusters,
						"This initializer can only be used on datasets of the same dimensions as the prior");
				return DirichletDistribution.sample(alpha, rnd);
			}
		};
	}
	
	/** Estimates the parameters (MAP) from an individual document for each cluster. **/
	public static POfXGivenYInitializer fromIndividualDocs(final double prior, final boolean includeLabeled,
			final RandomGenerator rnd) {
		return new POfXGivenYInitializer() {

			@Override
			public double[][] initialize(Dataset data, int numClusters) {
				Collection<SparseFeatureVector> instances;
				if (includeLabeled) {
					instances = Collections3.concat(Datasets.unlabel(data.labeledData()), data.unlabeledData());
				} else {
					instances = data.unlabeledData();
				}
				double [][] pOfXGivenY = new double[numClusters][];
				Iterable<SparseFeatureVector> samples =
						Collections3.sampleWithReplacement(instances, numClusters, rnd);
				for (Enumeration<SparseFeatureVector> e : Iterables2.enumerate(samples)) {
					pOfXGivenY[e.getIndex()] =
							Datasets.countFeatures(Collections.singleton(e.getElement()), data.getNumFeatures());
					DoubleArrays.addToSelf(pOfXGivenY[e.getIndex()], prior);
					DoubleArrays.normalizeToSelf(pOfXGivenY[e.getIndex()]);
				}
				return pOfXGivenY;
			}
			
		};
	}
			
	/** See Meila and Heckerman and Thiesson et al. (1999) **/
	public static POfXGivenYInitializer noisyMarginal(final double prior, final double effectiveSampleSize,
			final RandomGenerator rnd) {
		return new POfXGivenYInitializer() {

			@Override
			public double[][] initialize(Dataset data, int numClusters) {
				double[] empiricalPOfX = Datasets.countFeatures(data);
				DoubleArrays.normalizeToSelf(empiricalPOfX);
				DoubleArrays.multiplyToSelf(empiricalPOfX, effectiveSampleSize);
				// INFO : adding the prior after multiplying be the effectiveSampleSize avoids zeroes
				DoubleArrays.addToSelf(empiricalPOfX, prior);
				return DirichletDistribution.sample(new ArrayRealVector(empiricalPOfX, false), numClusters, rnd);
			}
			
		};
	}
	
	/**
	 * Combines the noisy marginal method with the posterior method in order to take into account labeled data
	 * when present; lambda is the amount of weight (in [0,1]) to give the posterior and effectiveSampleSize
	 * is the strength of the resulting Dirichlet. The parameters of the resulting Dirichlet have a mean equal to
	 * the the interpolation of the empirical p(x) and p(y|x).
	 **/
	public static POfXGivenYInitializer noisyMarginalInterpolatedWithPosterior(
			final double prior, final double lambda, final double effectiveSampleSize, final RandomGenerator rnd) {
		return new POfXGivenYInitializer() {

			@Override
			public double[][] initialize(Dataset data, int numClusters) {
				double[] empiricalPOfX = Datasets.countFeatures(data);
				DoubleArrays.normalizeToSelf(empiricalPOfX);
				
				double[][] empiricalPOfXGivenY = MixtureOfMultinomialsUtil.posteriorBeta(data, prior, numClusters);
				for (int k = 0; k < empiricalPOfXGivenY.length; k++) {
					DoubleArrays.normalizeToSelf(empiricalPOfXGivenY[k]);
					
					// Interpolate with empiricalPOfX
					DoubleArrays.addToSelfWeighted(empiricalPOfXGivenY[k], lambda, empiricalPOfX, 1.0 - lambda);

					// Set the effective sample size
					DoubleArrays.multiplyToSelf(empiricalPOfXGivenY[k], effectiveSampleSize);
					
					// Add the prior
					DoubleArrays.addToSelf(empiricalPOfXGivenY[k], prior);
				}
				return DirichletDistribution.sample(empiricalPOfXGivenY, rnd);
			}
			
		};
	}
	
	public static POfXGivenYInitializer sampleFromPosteriorPOfXGivenY(final double prior, final RandomGenerator rnd) {
		return new POfXGivenYInitializer() {
			
			@Override
			public double[][] initialize(Dataset data, int numClusters) {
				return DirichletDistribution.sample(MixtureOfMultinomialsUtil.posteriorBeta(data, prior, numClusters), rnd);
			}
		};
	}
	
}
